#自己定义的文件
import torch
import random
import json
import jsonlines
import time
import argparse
import sklearn


from torch.utils import data
from tqdm import tqdm
from apex import amp
from scipy.special import softmax

from ditto_light.ditto import evaluate, DittoModel
from ditto_light.exceptions import ModelNotFoundError
from ditto_light.dataset import DittoDataset
from ditto_light.summarize import Summarizer
from ditto_light.knowledge import *

def load_model(task, path, lm, use_gpu, fp16=True):
    checkpoint = os.path.join(path, task, 'model.pt')
    if not os.path.exists(checkpoint):
        raise ModelNotFoundError(checkpoint)

    configs = json.load(open('configs.json'))
    configs = {conf['name'] : conf for conf in configs}
    config = configs[task]
    config_list = [config]

    if use_gpu:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    else:
        device = 'cpu'

    model = DittoModel(device=device, lm=lm)

    saved_state = torch.load(checkpoint, map_location=lambda storage, loc: storage)
    model.load_state_dict(saved_state['model'])
    model.fingerprint = saved_state['tSNE']
    model = model.to(device)

    if fp16 and 'cuda' in device:
        model = amp.initialize(model, opt_level='O2')

    return config, model

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--lm", type=str, default='distilbert')
    parser.add_argument("--use_gpu", dest="use_gpu", action="store_true")
    parser.add_argument("--fp16", dest="fp16", action="store_true")
    parser.add_argument("--dk", type=str, default=None)
    parser.add_argument("--summarize", dest="summarize", action="store_true")
    parser.add_argument("--max_len", type=int, default=64)
    hp = parser.parse_args()
    task = 'wdc_watches_small' #选用一个现有源模型
    checkpoint_path = 'checkpoints/'
    lm = 'distilbert'
    use_gpu = True
    fp16 = True
    load = False
    if use_gpu:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    else:
        device = 'cpu'
    #加载训练之后的模型
    if load:
        _, model = load_model(task, checkpoint_path,
                               lm, use_gpu, fp16)
        #print(model.fingerprint)
    else:
        model = DittoModel(device=device, lm=lm)
        model = model.to(device)

    configs = json.load(open('configs.json'))
    configs = {conf['name']: conf for conf in configs}
    target_task_list = ["wdc_watches_small","wdc_shoes_small","wdc_computers_small",
                        "wdc_cameras_small","Structured/Walmart-Amazon",
                        "Structured/Amazon-Google","Structured/Beer", "Structured/DBLP-ACM",
                        "Structured/DBLP-GoogleScholar", "Structured/iTunes-Amazon"]
    res_list = []
    # 获取不同任务通过Bert捕获的特行
    for target_task in target_task_list:
        config = configs[target_task]
        dataset = config['trainset']
        train_dataset = DittoDataset(dataset,
                                     max_len=hp.max_len,#一般设置32或者64
                                     lm=hp.lm)
        iterator = data.DataLoader(dataset=train_dataset,
                                   batch_size=100,
                                   shuffle=True,#表示是否打乱
                                   num_workers=0,
                                   collate_fn=DittoDataset.pad)
        with torch.no_grad():
            # print('Classification')
            for i, batch in enumerate(iterator):
                batch1 = []
                for element in batch:
                    batch1.append(element.to(model.device))
                x, _ = batch1
                enc = model.feature(x)
                res = enc.cpu().numpy()

                res_list.append(res)
                break

    f=open("temp/enc_select.csv", "a")
    for res in res_list:
        np.savetxt(f,res,fmt='%.6f')
    f.close()